﻿#include "connect_socket.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <arpa/inet.h>
#include <sys/epoll.h>
#include <errno.h>
#include <pthread.h>

#include "log.h"




#define MAX_SOCKETS 1024
#define TIMEOUT_SEC 3  // epoll超时时间（秒）
#define BUFFER_SIZE 1024

typedef struct DataPtr {
    int tun_fd;
    HashMap* hashmap;
} DataPtr;

pthread_mutex_t sock_mutex;

int pipefd[2];  // 管道文件描述符
SocketNode pipeNode;

SocketNode rootNode;
SocketNode* endNode = &rootNode;


// 创建管道并添加到epoll
void create_pipe_and_add_to_epoll(int epoll_fd) {
    if (pipe(pipefd) == -1) {
        perror("pipe");
        exit(EXIT_FAILURE);
    }

    struct epoll_event event;
    event.events = EPOLLIN; // 监听读事件

    
    pipeNode.fd = pipefd[0];
    event.data.ptr = &pipeNode;

    if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, pipefd[0], &event) == -1) {
        perror("epoll_ctl: add pipe");
        close(pipefd[0]);
        close(pipefd[1]);
        exit(EXIT_FAILURE);
    }
}

// 唤醒epoll_wait
void wake_up_epoll() {
    char dummy = '1';
    write(pipefd[1], &dummy, 1);  // 写入管道以唤醒epoll_wait
    log_debug("唤醒 epoll");
}

SocketNode* create_sockets(char* src_ip, unsigned short int src_port, char* dst_ip, unsigned short int dst_port, unsigned  int ack) {
    struct sockaddr_in serv_addr;
    char* device = "eth0";  // 指定网络设备

    int sock = socket(AF_INET, SOCK_STREAM, 0);
    if (sock < 0) {
        perror("Socket creation failed");
        return NULL;
    }

    // 绑定套接字到指定网络设备
    if (setsockopt(sock, SOL_SOCKET, SO_BINDTODEVICE, device, (socklen_t)strlen(device) + 1) < 0) {
        perror("setsockopt failed");
        close(sock);
        return NULL;
    }

    // 设置套接字为非阻塞
    int flags = fcntl(sock, F_GETFL, 0);
    fcntl(sock, F_SETFL, flags | O_NONBLOCK);

    // 设置服务器地址
    serv_addr.sin_family = AF_INET;
    serv_addr.sin_port = htons(src_port);
    inet_pton(AF_INET, src_ip, &serv_addr.sin_addr);

    // 尝试连接
    int connect_result = connect(sock, (struct sockaddr*)&serv_addr, sizeof(serv_addr));
    if (connect_result < 0 && errno != EINPROGRESS) {
        if (errno == EINTR) {
            close(sock);
        }
        else {
            perror("Connection failed");
            close(sock);
        }
        return NULL;
    }

    // 成功创建并连接套接字
    pthread_mutex_lock(&sock_mutex);

    SocketNode* new_node = (SocketNode*)malloc(sizeof(SocketNode));
    if (new_node == NULL) {
        perror("malloc");
        return NULL;
    }
    new_node->fd = sock;
    new_node->next = NULL;
    new_node->seq = 1000;
    new_node->ack = ack;
    new_node->src_ip = src_ip;
    new_node->src_port = src_port;
    new_node->dst_ip = dst_ip;
    new_node->dst_port = dst_port;
    new_node->close = 0;
    pthread_mutex_init(&new_node->lock, NULL);


    // 动态分配内存来存储 src_ip 和 dst_ip
    new_node->src_ip = (char*)malloc(INET_ADDRSTRLEN );
    if (new_node->src_ip == NULL) {
        perror("malloc");
        free(new_node);
        close(sock);
        pthread_mutex_unlock(&sock_mutex);
        return NULL;
    }
    strncpy(new_node->src_ip, src_ip, INET_ADDRSTRLEN);
    new_node->src_ip[INET_ADDRSTRLEN - 1] = '\0';  // 确保字符串以 null 结尾

    new_node->dst_ip = (char*)malloc(INET_ADDRSTRLEN );
    if (new_node->dst_ip == NULL) {
        perror("malloc");
        free(new_node->src_ip);
        free(new_node);
        close(sock);
        pthread_mutex_unlock(&sock_mutex);
        return NULL;
    }
    strncpy(new_node->dst_ip, dst_ip, INET_ADDRSTRLEN);
    new_node->dst_ip[INET_ADDRSTRLEN - 1] = '\0';  // 确保字符串以 null 结尾


    endNode->next = new_node;
    endNode = new_node;

    pthread_mutex_unlock(&sock_mutex);

    log_info("创建socket描述符成功,fd=%d,[%s:%d-%s:%d]", sock, new_node->src_ip,new_node->src_port,new_node->dst_ip,new_node->dst_port);

    // 唤醒epoll_wait
    wake_up_epoll();

    return new_node;
}

void close_socket_node(int tun_fd,SocketNode* socket_node, HashMap* hashmap) {

    pthread_mutex_lock(&socket_node->lock);

    if (socket_node->close!=1) {
        socket_node->close = 1;
        close(socket_node->fd);
        log_info("关闭socket描述符，fn=%d", socket_node->fd);
        remove_socket_node(socket_node, hashmap);
        packet_fin(tun_fd, socket_node);
    }
    pthread_mutex_unlock(&socket_node->lock);

}

void remove_socket_node(SocketNode* socket_node,HashMap* hashmap) {
    char dest[30];
    int result = sprintf(dest, "%s#%d", socket_node->dst_ip, socket_node->dst_port);
    if (result>0) {
        hash_delete(hashmap,dest);
    }
}

void* monitor_sockets(void* arg) {

    DataPtr* data_ptr = ((DataPtr*)arg);
    int tun_fd = data_ptr->tun_fd;
    HashMap* hashmap = data_ptr->hashmap;

    free(arg);  // 释放分配的内存
    
    log_info("线程开始epoll监听\n");

    int epoll_fd = epoll_create1(0);
    if (epoll_fd == -1) {
        perror("epoll_create1");
        return NULL;
    }

    // 创建管道并添加到epoll
    create_pipe_and_add_to_epoll(epoll_fd);


    struct epoll_event event;
    event.events = EPOLLIN | EPOLLOUT | EPOLLERR | EPOLLHUP; // 监听读写和错误事件

    while (1) {

        pthread_mutex_lock(&sock_mutex);

        SocketNode* current = rootNode.next; 
        SocketNode* lastValidNode = &rootNode;     

        while (current != NULL) {
            SocketNode* next = current->next;  

            if (current->close != 1) {

                lastValidNode->next = current;

                event.data.ptr = current;
                if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, current->fd, &event) == -1 && errno != EEXIST) {
                    log_error("添加文件描述符fd=%d监听失败：%s", current->fd, strerror(errno));
                }
                else {
                    log_debug("添加文件描述符fd=%d监听成功。", current->fd);
                }

                // 更新最后一个有效节点
                lastValidNode = current;
            }
            else {
                //释放节点
                pthread_mutex_destroy(&current->lock);
                free(current->src_ip);
                free(current->dst_ip);
                free(current);
            }

            // 移动到下一个节点
            current = next;
        }

        // 如果有最后一个有效节点，确保其next指针为NULL
        if (lastValidNode != NULL && lastValidNode->next != NULL) {
            lastValidNode->next = NULL;
        }

        endNode = lastValidNode;

        pthread_mutex_unlock(&sock_mutex);

        struct epoll_event events[MAX_SOCKETS];
        int num_events = epoll_wait(epoll_fd, events, MAX_SOCKETS, TIMEOUT_SEC * 1000);
        if (num_events < 0) {
            if (errno == EINTR) {
                continue;  // 被中断，重试
            }
            log_debug("epoll无事件触发");
            break;
        }

        for (int i = 0; i < num_events; i++) {

            SocketNode* node = events[i].data.ptr; // 获取节点指针  
            if (node == NULL) {
                continue;
            }

            int fd = node->fd;
          
            if (fd == pipefd[0]) {
                char buffer[1024];
                ssize_t bytes_read = read(pipefd[0], buffer, sizeof(buffer) - 1);  // 读取管道数据
                if (bytes_read > 0) {
                    
                }
                
                continue;
            }

            if (events[i].events & (EPOLLERR | EPOLLHUP)) {
                log_error("描述符fd=%d错误", fd);
                events[i].data.ptr = NULL;
                close_socket_node(tun_fd,node, hashmap);

            }
            else if (events[i].events & EPOLLIN) {
                char buffer[1024];
                ssize_t bytes_received = recv(fd, buffer, sizeof(buffer), 0);
                if (bytes_received <= 0) {
                    if (bytes_received == 0) {
                        log_info("描述符fd=%d关闭\n", fd);
                    }
                    else {
                        if (errno == EAGAIN || errno == EWOULDBLOCK) {
                            continue;  // 资源暂时不可用，重试
                        }
                        log_error("描述符fd=%d读取错误", strerror(errno));
                    }   
                    events[i].data.ptr = NULL;
                    close_socket_node(tun_fd,node, hashmap);

                }
                else {

                    packet_push(tun_fd,node,buffer,(size_t) bytes_received);
                    node->seq += (unsigned int)bytes_received;

                }
            }
            else if (events[i].events & EPOLLOUT) {
                log_debug("描述符fd=%d连接成功",fd);
                // 修改事件类型为读事件
                events[i].events = EPOLLIN | EPOLLERR | EPOLLHUP;

                if (epoll_ctl(epoll_fd, EPOLL_CTL_MOD, fd, &events[i]) == -1) {
                    perror("Epoll modify failed");
                    close(fd);
                }

                // 回应 syn+ack
                packet_syn_ack(tun_fd, node);
                node->seq += 1;

            }

        }

    }

    close(epoll_fd);
    return NULL;
}

void start_sockets(int tun_fd,HashMap * hashmap) {

    // 分配内存来存储 fd
    DataPtr* data_ptr = malloc(sizeof(DataPtr));
    if (data_ptr == NULL) {
        perror("Memory allocation failed");
        return;
    }
    data_ptr->tun_fd = tun_fd;
    data_ptr->hashmap = hashmap;

    pthread_t thread1;
    pthread_mutex_init(&sock_mutex, NULL);


    if (pthread_create(&thread1, NULL, monitor_sockets, data_ptr) != 0) {
        perror("Failed to create thread2");
        return ;
    }

}